{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "import cv2\n",
    "import importlib\n",
    "import torch\n",
    "import numpy as np\n",
    "import argparse\n",
    "import yaml\n",
    "from tqdm import tqdm\n",
    "from PIL import Image, ImageDraw, ImageOps\n",
    "import json\n",
    "import nltk\n",
    "import torchvision.transforms as T\n",
    "import matplotlib.pyplot as plt\n",
    "import torch.nn.functional as F\n",
    "from scipy import signal\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.widgets import Button\n",
    "import clip\n",
    "import kornia\n",
    "import torchvision\n",
    "import skimage.color\n",
    "import random\n",
    "\n",
    "from utils_func import *\n",
    "from html_images import *\n",
    "from sample_func import *\n",
    "from ImageMatch.warp import ImageWarper\n",
    "from colorizer import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create colorizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ckpt_file = 'C:/MyFiles/CondTran/finals/bert_final/logs/bert/epoch=14-step=142124.ckpt'\n",
    "device = 'cuda:0'\n",
    "colorizer = Colorizer(ckpt_file, device, [256, 256])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sample dataset (raw)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sampling arguments\n",
    "img_dir = 'C:\\\\MyFiles\\\\Dataset\\\\imagenet\\\\val5000\\\\val'\n",
    "save_dir = 'C:\\\\Users\\\\lucky\\\\Desktop\\\\raw_diverse'\n",
    "topk = 100\n",
    "num_samples = 5\n",
    "img_size = [256, 256]\n",
    "sample_size = [img_size[0]//16, img_size[1]//16]\n",
    "\n",
    "html = HTML(save_dir, 'Sample')\n",
    "\n",
    "files = os.listdir(img_dir)\n",
    "np.random.shuffle(files)\n",
    "\n",
    "pbar = tqdm(enumerate(files))\n",
    "for i, filename in pbar:\n",
    "    if filename.endswith('.jpg') or filename.endswith('.png') or filename.endswith('.JPEG') or filename.endswith('.jpeg'):\n",
    "        fname = filename.split('.')[0]\n",
    "        I_color = Image.open(os.path.join(img_dir, filename)).convert('RGB')\n",
    "        I_gray = I_color.convert('L')\n",
    "        \n",
    "        gen_imgs = [I_color, I_gray]\n",
    "\n",
    "        gen = colorizer.sample(I_gray, strokes=[], topk=topk)\n",
    "        gen_imgs.append(gen)\n",
    "        \n",
    "        save_result(html, index=fname, images=gen_imgs, texts=[fname])\n",
    "        html.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sample strokes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sampling arguments\n",
    "topk = 100\n",
    "num_samples = 1\n",
    "img_size = [256, 256]\n",
    "sample_size = [img_size[0]//16, img_size[1]//16]\n",
    "stroke_path = 'C:\\\\MyFiles\\\\CondTran\\\\data\\\\coco_strokes.json'\n",
    "dataset_path = 'C:\\\\MyFiles\\\\Dataset\\\\coco\\\\val2017'\n",
    "save_dir = 'C:\\\\Users\\\\lucky\\\\Desktop\\\\final_stroke_coco'\n",
    "num_strokes = [2, 16]\n",
    "\n",
    "html = HTML(save_dir, 'Sample')\n",
    "\n",
    "# Load strokes\n",
    "with open(stroke_path, 'r') as file:\n",
    "    all_strokes = json.load(file)\n",
    "    print(len(all_strokes))\n",
    "\n",
    "pbar = tqdm(enumerate(all_strokes))\n",
    "\n",
    "random.seed(100)\n",
    "for i, file in pbar:\n",
    "    filename = file['image']\n",
    "    strokes = file['strokes']\n",
    "    n_strokes = random.randint(num_strokes[0], num_strokes[1])\n",
    "    n_strokes = min(n_strokes, len(strokes))\n",
    "    strokes = random.sample(strokes, k=n_strokes)\n",
    "    name = filename.split('.')[0]\n",
    "\n",
    "    gen_imgs = []\n",
    "\n",
    "    I_color = Image.open(os.path.join(dataset_path, filename)).convert('RGB')\n",
    "    I_gray = I_color.convert('L')\n",
    "\n",
    "    gen_imgs.append(I_color)\n",
    "\n",
    "    draw_img = I_gray.copy().resize(img_size).convert('RGB')\n",
    "    draw_img = draw_strokes(draw_img, img_size, strokes)\n",
    "\n",
    "    gen_imgs.append(draw_img.resize(I_color.size))\n",
    "\n",
    "    for n in range(num_samples):\n",
    "        gen = colorizer.sample(I_gray, strokes, topk)\n",
    "        gen_imgs.append(gen)\n",
    "    \n",
    "    save_result(html, index=i, images=gen_imgs, texts=['original', 'strokes', 'colorized'])\n",
    "    html.save()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sample text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sampling arguments\n",
    "topk = 100\n",
    "num_samples = 1\n",
    "img_size = [256, 256]\n",
    "sample_size = [img_size[0]//16, img_size[1]//16]\n",
    "stroke_path = 'C:/MyFiles/CondTran/data/all_text_strokes.json'\n",
    "dataset_path = 'C:/MyFiles/Dataset/coco/val2017'\n",
    "save_dir = 'C:/MyFiles/CondTran/sample_result/slic_text'\n",
    "\n",
    "html = HTML(save_dir, 'Sample')\n",
    "html.add_header(os.path.join(args.root_dir, args.log_dir, args.step))\n",
    "\n",
    "# Load strokes\n",
    "with open(stroke_path, 'r') as file:\n",
    "    all_strokes = json.load(file)\n",
    "\n",
    "np.random.seed(10)\n",
    "np.random.shuffle(all_strokes)\n",
    "pbar = tqdm(enumerate(all_strokes[:200]))\n",
    "\n",
    "for i, file in pbar:\n",
    "    filename = file['image']\n",
    "    strokes = file['strokes']\n",
    "    name = filename.split('.')[0]\n",
    "\n",
    "    for stk in strokes:\n",
    "        ind = stk['index'].copy()\n",
    "        stk['index'] = [ind[0]*16, ind[1]*16]\n",
    "\n",
    "    gen_imgs = []\n",
    "\n",
    "    I_color = Image.open(os.path.join(dataset_path, filename)).convert('RGB')\n",
    "    I_gray = I_color.convert('L')\n",
    "\n",
    "    gen_imgs.append(I_color)\n",
    "\n",
    "    draw_img = I_gray.copy().resize(img_size).convert('RGB')\n",
    "    for stk in strokes:\n",
    "        ind = stk['index']\n",
    "        color = np.array(stk['color'])\n",
    "        color = np.expand_dims(color, axis=(0, 1))\n",
    "        color = cv2.resize(color, (16-6, 16-6), interpolation=cv2.INTER_NEAREST)\n",
    "        color = cv2.copyMakeBorder(color, 3, 3, 3, 3, cv2.BORDER_CONSTANT, value=(255, 255, 255))\n",
    "        draw_img = draw_full_color(draw_img, color, [ind[0], ind[0]+16, ind[1], ind[1]+16])\n",
    "\n",
    "    gen_imgs.append(draw_img.resize(I_color.size))\n",
    "    \n",
    "    x_color = preprocess(I_color, img_size).to(args.device)\n",
    "    x_gray = preprocess(I_gray, img_size).to(args.device)\n",
    "\n",
    "    for n in range(num_samples):\n",
    "        gen = filltran.sample(x_gray, topk, strokes)\n",
    "        gen = output_to_pil(gen[0])\n",
    "        gen_resize = color_resize(I_gray, gen)\n",
    "        gen_imgs.append(gen_resize)\n",
    "    \n",
    "    save_result(html, index=i, images=gen_imgs, texts=[' ', file['caption']])\n",
    "    html.save()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Sample exemplar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sampling arguments\n",
    "topk = 100\n",
    "num_samples = 1\n",
    "img_size = [256, 256]\n",
    "sample_size = [img_size[0]//16, img_size[1]//16]\n",
    "pair_path = 'C:/MyFiles/CondTran/data/all_exemplars.json'\n",
    "dataset_path = 'C:/MyFiles/Dataset/imagenet/val5000/val'\n",
    "ref_path = 'C:/MyFiles/Dataset/imagenet/full/train'\n",
    "save_dir = 'C:/MyFiles/CondTran/sample_result/bert_final_exp'\n",
    "\n",
    "html = HTML(save_dir, 'Sample')\n",
    "html.add_header(os.path.join(args.root_dir, args.log_dir, args.step))\n",
    "\n",
    "# Load strokes\n",
    "with open(pair_path, 'r') as file:\n",
    "    all_pairs = json.load(file)\n",
    "np.random.shuffle(all_pairs)\n",
    "\n",
    "# Load image warper\n",
    "warper = ImageWarper('cuda')\n",
    "\n",
    "i = 0\n",
    "pbar = tqdm(all_pairs)\n",
    "for file in pbar:\n",
    "    filename = file['image']\n",
    "    refname = file['exemplar']\n",
    "    name = filename.split('.')[0]\n",
    "    gen_imgs = []\n",
    "\n",
    "    in_dir = os.path.join(dataset_path, filename)\n",
    "    ref_dir = os.path.join(ref_path, refname)\n",
    "    I_color = Image.open(in_dir).convert('RGB')\n",
    "    I_gray = I_color.convert('L')\n",
    "    I_ref = Image.open(ref_dir).convert('RGB')\n",
    "\n",
    "    gen_imgs.append(I_color)\n",
    "    gen_imgs.append(I_ref)\n",
    "\n",
    "    warped_img, similarity_map = warper.warp_image(I_gray.convert('RGB'), I_ref)\n",
    "    gen_imgs.append(warped_img.resize(I_color.size))\n",
    "    warped_img = warped_img.resize(img_size)\n",
    "\n",
    "    similarity_map = cv2.resize(similarity_map, tuple(sample_size))\n",
    "    similarity_map = similarity_map.reshape(-1)\n",
    "    threshold = min(0.23, np.sort(similarity_map)[-10])\n",
    "    indices = np.where( (similarity_map >= 0.23))\n",
    "\n",
    "    strokes = []\n",
    "    warped_img = np.array(warped_img)\n",
    "    for ind in indices[0]:\n",
    "        index = [ind//16 * 16, ind%16 * 16]\n",
    "        color = warped_img[index[0]:index[0]+16, index[1]:index[1]+16, :]\n",
    "        color = color.mean(axis=(0, 1))\n",
    "        strokes.append({'index': index, 'color': color.tolist()})\n",
    "\n",
    "    draw_img = I_gray.copy().resize(img_size).convert('RGB')\n",
    "    for stk in strokes:\n",
    "        ind = stk['index']\n",
    "        color = np.array(stk['color'])\n",
    "        color = np.expand_dims(color, axis=(0, 1))\n",
    "        color = cv2.resize(color, (16-6, 16-6), interpolation=cv2.INTER_NEAREST)\n",
    "        color = cv2.copyMakeBorder(color, 3, 3, 3, 3, cv2.BORDER_CONSTANT, value=(255, 255, 255))\n",
    "        draw_img = draw_full_color(draw_img, color, [ind[0], ind[0]+16, ind[1], ind[1]+16])\n",
    "\n",
    "    gen_imgs.append(draw_img.resize(I_color.size))\n",
    "\n",
    "    x_gray = preprocess(I_gray, img_size).to(args.device)\n",
    "\n",
    "    for n in range(num_samples):\n",
    "        gen = filltran.sample(x_gray, topk, strokes)\n",
    "        gen = output_to_pil(gen[0])\n",
    "        #gen.show()\n",
    "        gen_resize = color_resize(I_gray, gen)\n",
    "        #gen_resize.show()\n",
    "        gen_imgs.append(gen_resize)\n",
    "\n",
    "    save_result(html, index=i, images=gen_imgs)\n",
    "    html.save()\n",
    "    i += 1\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Upsample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sampling arguments\n",
    "img_dir = 'C:/MyFiles/Dataset/imagenet/val5000/val'\n",
    "save_dir = 'C:/MyFiles/CondTran/sample_result/bert_final_upsample'\n",
    "topk = 1\n",
    "img_size = [256, 256]\n",
    "\n",
    "html = HTML(save_dir, 'Sample')\n",
    "\n",
    "pbar = tqdm(enumerate(os.listdir(img_dir)[:10]))\n",
    "for i, filename in pbar:\n",
    "    if filename.endswith('.jpg') or filename.endswith('.png') or filename.endswith('.JPEG') or filename.endswith('.jpeg'):\n",
    "        gen_imgs = []\n",
    "\n",
    "        I_color = Image.open(os.path.join(img_dir, filename)).convert('RGB')\n",
    "        gen_imgs.append(I_color)\n",
    "        I_gray = I_color.convert('L')\n",
    "        I_color = I_color.resize(img_size)\n",
    "        \n",
    "        I_bilinear = color_resize(I_gray, I_color)\n",
    "        gen_imgs.append(I_bilinear)\n",
    "\n",
    "        I_upsample = colorizer.upsample(I_gray, I_color)\n",
    "        gen_imgs.append(I_upsample)\n",
    "\n",
    "        save_result(html, index=i, images=gen_imgs)\n",
    "        html.save()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.3 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "d2a91ca3d26e2d2b798fe8778458ed9052c39446826a52af3c736ec2371151af"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
